{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Controlling Agent Reasoning Loop with Return Direct Tools\n",
    "\n",
    "All tools have an option for `return_direct` -- if this is set to `True`, and the associated tool is called (without any other tools being called), the agent reasoning loop is ended and the tool output is returned directly.\n",
    "\n",
    "This can be useful for speeding up response times when you know the tool output is good enough, to avoid the agent re-writing the response, and for ending the reasoning loop."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This notebook walks through a notebook where an agent needs to gather information from a user in order to make a restaurant booking."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%pip install llama-index-core llama-index-llms-anthropic"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "\n",
    "os.environ[\"ANTHROPIC_API_KEY\"] = \"sk-...\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Tools setup"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import Optional\n",
    "\n",
    "from llama_index.core.tools import FunctionTool\n",
    "from pydantic import BaseModel\n",
    "\n",
    "# we will store booking under random IDs\n",
    "bookings = {}\n",
    "\n",
    "\n",
    "# we will represent and track the state of a booking as a Pydantic model\n",
    "class Booking(BaseModel):\n",
    "    name: Optional[str] = None\n",
    "    email: Optional[str] = None\n",
    "    phone: Optional[str] = None\n",
    "    date: Optional[str] = None\n",
    "    time: Optional[str] = None\n",
    "\n",
    "\n",
    "def get_booking_state(user_id: str) -> str:\n",
    "    \"\"\"Get the current state of a booking for a given booking ID.\"\"\"\n",
    "    try:\n",
    "        return str(bookings[user_id].dict())\n",
    "    except:\n",
    "        return f\"Booking ID {user_id} not found\"\n",
    "\n",
    "\n",
    "def update_booking(user_id: str, property: str, value: str) -> str:\n",
    "    \"\"\"Update a property of a booking for a given booking ID. Only enter details that are explicitly provided.\"\"\"\n",
    "    booking = bookings[user_id]\n",
    "    setattr(booking, property, value)\n",
    "    return f\"Booking ID {user_id} updated with {property} = {value}\"\n",
    "\n",
    "\n",
    "def create_booking(user_id: str) -> str:\n",
    "    \"\"\"Create a new booking and return the booking ID.\"\"\"\n",
    "    bookings[user_id] = Booking()\n",
    "    return \"Booking created, but not yet confirmed. Please provide your name, email, phone, date, and time.\"\n",
    "\n",
    "\n",
    "def confirm_booking(user_id: str) -> str:\n",
    "    \"\"\"Confirm a booking for a given booking ID.\"\"\"\n",
    "    booking = bookings[user_id]\n",
    "\n",
    "    if booking.name is None:\n",
    "        raise ValueError(\"Please provide your name.\")\n",
    "\n",
    "    if booking.email is None:\n",
    "        raise ValueError(\"Please provide your email.\")\n",
    "\n",
    "    if booking.phone is None:\n",
    "        raise ValueError(\"Please provide your phone number.\")\n",
    "\n",
    "    if booking.date is None:\n",
    "        raise ValueError(\"Please provide the date of your booking.\")\n",
    "\n",
    "    if booking.time is None:\n",
    "        raise ValueError(\"Please provide the time of your booking.\")\n",
    "\n",
    "    return f\"Booking ID {user_id} confirmed!\"\n",
    "\n",
    "\n",
    "# create tools for each function\n",
    "get_booking_state_tool = FunctionTool.from_defaults(fn=get_booking_state)\n",
    "update_booking_tool = FunctionTool.from_defaults(fn=update_booking)\n",
    "create_booking_tool = FunctionTool.from_defaults(\n",
    "    fn=create_booking, return_direct=True\n",
    ")\n",
    "confirm_booking_tool = FunctionTool.from_defaults(\n",
    "    fn=confirm_booking, return_direct=True\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## A user has walked in! Let's help them make a booking"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from llama_index.llms.anthropic import Anthropic\n",
    "from llama_index.core.llms import ChatMessage\n",
    "from llama_index.core.agent.workflow import FunctionAgent\n",
    "from llama_index.core.workflow import Context\n",
    "\n",
    "llm = Anthropic(model=\"claude-3-sonnet-20240229\", temperature=0.1)\n",
    "\n",
    "user = \"user123\"\n",
    "system_prompt = f\"\"\"You are now connected to the booking system and helping {user} with making a booking.\n",
    "Only enter details that the user has explicitly provided.\n",
    "Do not make up any details.\n",
    "\"\"\"\n",
    "\n",
    "agent = FunctionAgent(\n",
    "    tools=[\n",
    "        get_booking_state_tool,\n",
    "        update_booking_tool,\n",
    "        create_booking_tool,\n",
    "        confirm_booking_tool,\n",
    "    ],\n",
    "    llm=llm,\n",
    "    system_prompt=system_prompt,\n",
    ")\n",
    "\n",
    "# create a context for the agent to hold the state/history of a session\n",
    "ctx = Context(agent)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Okay, let's create a new booking for you.{\"user_id\": \"user123\"}\n",
      "Call create_booking with {'user_id': 'user123'}\n",
      "Returned: Booking created, but not yet confirmed. Please provide your name, email, phone, date, and time.\n"
     ]
    }
   ],
   "source": [
    "from llama_index.core.agent.workflow import AgentStream, ToolCallResult\n",
    "\n",
    "handler = agent.run(\n",
    "    \"Hello! I would like to make a booking, around 5pm?\", ctx=ctx\n",
    ")\n",
    "\n",
    "async for ev in handler.stream_events():\n",
    "    if isinstance(ev, AgentStream):\n",
    "        print(f\"{ev.delta}\", end=\"\", flush=True)\n",
    "    elif isinstance(ev, ToolCallResult):\n",
    "        print(\n",
    "            f\"\\nCall {ev.tool_name} with {ev.tool_kwargs}\\nReturned: {ev.tool_output}\"\n",
    "        )\n",
    "\n",
    "response = await handler"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Booking created, but not yet confirmed. Please provide your name, email, phone, date, and time.\n"
     ]
    }
   ],
   "source": [
    "print(str(response))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Perfect, we can see the function output was retruned directly, with no modification or final LLM call!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Got it, thanks for providing your name and email. I've updated the booking with that information.{\"user_id\": \"user123\", \"property\": \"name\", \"value\": \"Logan\"}{\"user_id\": \"user123\", \"property\": \"email\", \"value\": \"test@gmail.com\"}\n",
      "Call update_booking with {'user_id': 'user123', 'property': 'name', 'value': 'Logan'}\n",
      "Returned: Booking ID user123 updated with name = Logan\n",
      "\n",
      "Call update_booking with {'user_id': 'user123', 'property': 'email', 'value': 'test@gmail.com'}\n",
      "Returned: Booking ID user123 updated with email = test@gmail.com\n",
      "Please also provide your phone number, preferred date, and time for the booking."
     ]
    }
   ],
   "source": [
    "handler = agent.run(\n",
    "    \"Sure! My name is Logan, and my email is test@gmail.com?\", ctx=ctx\n",
    ")\n",
    "\n",
    "async for ev in handler.stream_events():\n",
    "    if isinstance(ev, AgentStream):\n",
    "        print(f\"{ev.delta}\", end=\"\", flush=True)\n",
    "    elif isinstance(ev, ToolCallResult):\n",
    "        print(\n",
    "            f\"\\nCall {ev.tool_name} with {ev.tool_kwargs}\\nReturned: {ev.tool_output}\"\n",
    "        )\n",
    "\n",
    "response = await handler"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Please also provide your phone number, preferred date, and time for the booking.\n"
     ]
    }
   ],
   "source": [
    "print(str(response))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Great, thank you for providing the additional details. I've updated the booking with your phone number, date, and time.{\"user_id\": \"user123\", \"property\": \"phone\", \"value\": \"1234567890\"}{\"user_id\": \"user123\", \"property\": \"date\", \"value\": \"2023-04-05\"}{\"user_id\": \"user123\", \"property\": \"time\", \"value\": \"17:00\"}\n",
      "Call update_booking with {'user_id': 'user123', 'property': 'phone', 'value': '1234567890'}\n",
      "Returned: Booking ID user123 updated with phone = 1234567890\n",
      "\n",
      "Call update_booking with {'user_id': 'user123', 'property': 'date', 'value': '2023-04-05'}\n",
      "Returned: Booking ID user123 updated with date = 2023-04-05\n",
      "\n",
      "Call update_booking with {'user_id': 'user123', 'property': 'time', 'value': '17:00'}\n",
      "Returned: Booking ID user123 updated with time = 17:00\n",
      "Looks like I have all the necessary details. Let me confirm this booking for you.{\"user_id\": \"user123\"}\n",
      "Call confirm_booking with {'user_id': 'user123'}\n",
      "Returned: Booking ID user123 confirmed!\n"
     ]
    }
   ],
   "source": [
    "handler = agent.run(\n",
    "    \"Right! My phone number is 1234567890, the date of the booking is April 5, at 5pm.\",\n",
    "    ctx=ctx,\n",
    ")\n",
    "\n",
    "async for ev in handler.stream_events():\n",
    "    if isinstance(ev, AgentStream):\n",
    "        print(f\"{ev.delta}\", end=\"\", flush=True)\n",
    "    elif isinstance(ev, ToolCallResult):\n",
    "        print(\n",
    "            f\"\\nCall {ev.tool_name} with {ev.tool_kwargs}\\nReturned: {ev.tool_output}\"\n",
    "        )\n",
    "\n",
    "response = await handler"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Booking ID user123 confirmed!\n"
     ]
    }
   ],
   "source": [
    "print(str(response))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "name='Logan' email='test@gmail.com' phone='1234567890' date='2023-04-05' time='17:00'\n"
     ]
    }
   ],
   "source": [
    "print(bookings[\"user123\"])"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "venv",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
